import numpy
import interpolation

def lightcurve(file,
               interpolant = interpolation.least_squares_polynomial,
               evaluator = interpolation.polynomial_evaluator,
               degree = 5,
               min_period = 0.2,
               max_period = 32.,
               period_bins = 50000,
               sigma = 3,
               min_obs = 25):
    """Returns a four-tuple containing 
    1) the name of the star, 
    2) the period from a Lomb-Scargle Periodogram, 
    3) the rephased observations, and 
    4) a list of interpolation coefficients, or 
    None, None, None, None if no adequate model can be found. 
    
    Searches for periods within the specified bounds discretized by the 
    specified number of period bins. 
    
    If sigma is greater than zero, then outliers will be sigma clipped with the 
    specified value. 
    """
    from os import sep
    from re import split
    name = file.split(sep)[-1]
    data = numpy.ma.masked_array(data=numpy.loadtxt(file), mask=None)
    failure = None, None, None, None
    while True:
        if get_unmasked(data).shape[0] < min_obs:
            return failure
        period = find_period(data, min_period, max_period, period_bins)
        if not min_period <= period <= max_period:
            return failure
        rephased = rephase(data, period)
        coefficients = interpolant(rephased, degree)
        if sigma <= 0:
            return name, period, rephased, coefficients
        prev_mask = data.mask
        outliers = find_outliers(rephased, evaluator, coefficients, sigma)
        data.mask = numpy.ma.mask_or(data.mask, outliers)
        if numpy.all(data.mask == prev_mask):
            return name, period, rephased, coefficients

def get_unmasked(data):
    """Returns all of the values in a specified column that are not outliers."""
    return data[~data.mask].reshape(-1, data.shape[1])

def find_period(data, min_period, max_period, period_bins):
    """Uses the Lomb-Scargle Periodogram to discover the period."""
    from scipy.signal import lombscargle
    time, mags = data.T[0], data.T[1]
    scaled_mags = (mags-mags.mean())/mags.std()
    freqs = numpy.linspace(2*numpy.pi/max_period, 2*numpy.pi/min_period, period_bins)
    pgram = lombscargle(time, scaled_mags, freqs)
    return 2*numpy.pi/freqs[numpy.argmax(pgram)]

def rephase(data, period):
    """Non-destructively rephases all of the points in the given data set to
    be between 0 and 1 and shifted to max light."""
    max_light_phase = get_phase(max_light_time(data), period)
    rephased = numpy.ma.copy(data)
    for observation in rephased:
        observation[0] = get_phase(observation[0], period, offset=max_light_phase)
    return rephased

def max_light_time(data):
    """Returns the time at which the star is at its brightest, i.e. the Julian
    Date of the smallest magnitude."""
    return data.T[0][data.T[1].argmin()]

def get_phase(time, period, offset=0):
    """Returns the phase associated with a given time based on the period."""
    from math import modf
    return (modf(time/period)[0]-offset)%1

def find_outliers(rephased, evaluator, coefficients, sigma):
    """Finds rephased values that are too far from the light curve."""
    expected = evaluator(rephased.T[0], coefficients)
    actual, error = rephased.T[1], rephased.T[2]
    clip = sigma*actual.std()**2
    outliers = (expected-actual)**2 > clip+error
    return numpy.tile(numpy.vstack(outliers), rephased.shape[1])